{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notebook written by [Zhedong Zheng](https://github.com/zhedongzheng)\n",
    "\n",
    "<img src=\"img/pointer_net.png\" width=\"500\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "PARAMS = {\n",
    "    'max_len': 15,\n",
    "    'embed_dims': 15,\n",
    "    'rnn_size': 50,\n",
    "    'clip_norm': 5.0,\n",
    "    'batch_size': 128,\n",
    "    'n_epochs': 100,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_data(path):\n",
    "    with open(path, 'r', encoding='utf-8') as f:\n",
    "        return f.read()\n",
    "\n",
    "    \n",
    "def build_map(data):\n",
    "    specials = ['<PAD>', '<GO>',  '<EOS>', '<UNK>']\n",
    "    chars = list(set([char for line in data.split('\\n') for char in line]))\n",
    "    chars = sorted(chars)\n",
    "    idx2char = {idx: char for idx, char in enumerate(specials + chars)}\n",
    "    char2idx = {char: idx for idx, char in idx2char.items()}\n",
    "    return idx2char, char2idx\n",
    "\n",
    "\n",
    "def preprocess_data():\n",
    "    source = read_data('../temp/letters_source.txt')\n",
    "    target = read_data('../temp/letters_target.txt')\n",
    "\n",
    "    PARAMS['src_idx2char'], PARAMS['src_char2idx'] = build_map(source)\n",
    "    \n",
    "    src_indices, tgt_indices = [], []\n",
    "    src_seq_lens, tgt_seq_lens = [], []\n",
    "    \n",
    "    for src_line, tgt_line in zip(source.split('\\n'), target.split('\\n')):\n",
    "        src_idx = [PARAMS['src_char2idx'].get(c, 3) for c in src_line] + [2]\n",
    "        src_seq_lens.append(len(src_idx))\n",
    "        src_idx = src_idx + [0] * (PARAMS['max_len']-len(src_idx))\n",
    "        \n",
    "        tgt_idx = [PARAMS['src_char2idx'].get(c, 3) for c in tgt_line] + [2]\n",
    "        tgt_seq_lens.append(len(tgt_idx))\n",
    "        tgt_idx = tgt_idx + [0] * (PARAMS['max_len']-len(tgt_idx))\n",
    "        tgt_idx = [src_idx.index(t) for t in tgt_idx]\n",
    "        \n",
    "        src_indices.append(src_idx)\n",
    "        tgt_indices.append(tgt_idx)\n",
    "    \n",
    "    return (np.array(src_indices),\n",
    "            np.array(tgt_indices),\n",
    "            np.array(src_seq_lens),\n",
    "            np.array(tgt_seq_lens))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def clip_grads(loss):\n",
    "    variables = tf.trainable_variables()\n",
    "    grads = tf.gradients(loss, variables)\n",
    "    clipped_grads, _ = tf.clip_by_global_norm(grads, PARAMS['clip_norm'])\n",
    "    return zip(clipped_grads, variables)\n",
    "\n",
    "\n",
    "def rnn_cell():\n",
    "    return tf.nn.rnn_cell.GRUCell(PARAMS['rnn_size'],\n",
    "                                  kernel_initializer=tf.orthogonal_initializer())\n",
    "\n",
    "\n",
    "def point(idx, batch_sz, enc_inp):\n",
    "    return tf.gather_nd(enc_inp, tf.concat([\n",
    "        tf.expand_dims(tf.range(batch_sz), 1),\n",
    "        tf.expand_dims(idx, 1)],\n",
    "        axis=1))\n",
    "\n",
    "\n",
    "def attention(query, keys, masks, W1, W2, v):\n",
    "    query = tf.expand_dims(query, 1)\n",
    "    align = v * tf.tanh(W1(query) + W2(keys))\n",
    "    align = tf.reduce_sum(align, [2])\n",
    "    align *= masks\n",
    "    return align\n",
    "\n",
    "\n",
    "def forward(features):\n",
    "    inputs = features['src_idx']\n",
    "    enc_seq_len = features['src_seq_lens']\n",
    "    batch_sz = tf.shape(inputs)[0]\n",
    "    masks = tf.to_float(tf.sign(inputs))\n",
    "    \n",
    "    with tf.variable_scope('Encoder'):\n",
    "        embedding = tf.get_variable('lookup_table',\n",
    "                                    [len(PARAMS['src_char2idx']), PARAMS['embed_dims']])\n",
    "        enc_inp = tf.nn.embedding_lookup(embedding, inputs)\n",
    "        enc_rnn_out, enc_rnn_state = tf.nn.dynamic_rnn(rnn_cell(),\n",
    "                                                       enc_inp,\n",
    "                                                       enc_seq_len,\n",
    "                                                       dtype=tf.float32)\n",
    "        \n",
    "    with tf.variable_scope('Decoder'):\n",
    "        outputs = []\n",
    "        \n",
    "        dec_cell = rnn_cell()\n",
    "        W1 = tf.layers.Dense(PARAMS['rnn_size'], use_bias=False)\n",
    "        W2 = tf.layers.Dense(PARAMS['rnn_size'], use_bias=False)\n",
    "        v = tf.get_variable('v', [PARAMS['rnn_size']])\n",
    "        \n",
    "        state = enc_rnn_state\n",
    "        starts = tf.fill([batch_sz], PARAMS['src_char2idx']['<GO>'])\n",
    "        inp = tf.nn.embedding_lookup(embedding, starts)\n",
    "        \n",
    "        for _ in range(PARAMS['max_len']):\n",
    "            _, state = dec_cell(inp, state)\n",
    "            output = attention(state, enc_rnn_out, masks, W1, W2, v)\n",
    "            outputs.append(output)\n",
    "            idx = tf.argmax(output, -1, output_type=tf.int32)\n",
    "            inp = point(idx, batch_sz, enc_inp)\n",
    "    \n",
    "    outputs = tf.stack(outputs, 1)\n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_fn(features, labels, mode, params):\n",
    "    logits = forward(features)\n",
    "    \n",
    "    if mode == tf.estimator.ModeKeys.PREDICT:\n",
    "        return tf.estimator.EstimatorSpec(mode, predictions=tf.argmax(logits, -1))\n",
    "        \n",
    "    if mode == tf.estimator.ModeKeys.TRAIN:\n",
    "        loss_op = tf.contrib.seq2seq.sequence_loss(\n",
    "            logits = logits,\n",
    "            targets = labels['tgt_idx'],\n",
    "            weights = tf.sequence_mask(labels['tgt_seq_lens'], PARAMS['max_len'], dtype=tf.float32))\n",
    "        train_op = tf.train.AdamOptimizer().apply_gradients(\n",
    "            clip_grads(loss_op),\n",
    "            global_step = tf.train.get_global_step())\n",
    "        \n",
    "        return tf.estimator.EstimatorSpec(mode=mode, loss=loss_op, train_op=train_op)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using default config.\n",
      "WARNING:tensorflow:Using temporary folder as model directory: /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpxxtascq7\n",
      "INFO:tensorflow:Using config: {'_model_dir': '/var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpxxtascq7', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': None, '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': None, '_service': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x12187fa90>, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1}\n",
      "WARNING:tensorflow:Estimator's model_fn (<function model_fn at 0x12183c9d8>) includes params argument, but params are not passed to Estimator.\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Create CheckpointSaverHook.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "INFO:tensorflow:Saving checkpoints for 1 into /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpxxtascq7/model.ckpt.\n",
      "INFO:tensorflow:loss = 2.736074, step = 1\n",
      "INFO:tensorflow:global_step/sec: 14.9557\n",
      "INFO:tensorflow:loss = 1.6095542, step = 101 (6.687 sec)\n",
      "INFO:tensorflow:global_step/sec: 20.814\n",
      "INFO:tensorflow:loss = 1.2738905, step = 201 (4.804 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.7768\n",
      "INFO:tensorflow:loss = 1.0263916, step = 301 (4.206 sec)\n",
      "INFO:tensorflow:global_step/sec: 20.4712\n",
      "INFO:tensorflow:loss = 0.78070325, step = 401 (4.885 sec)\n",
      "INFO:tensorflow:global_step/sec: 21.782\n",
      "INFO:tensorflow:loss = 0.5283796, step = 501 (4.591 sec)\n",
      "INFO:tensorflow:global_step/sec: 20.7859\n",
      "INFO:tensorflow:loss = 0.4521801, step = 601 (4.811 sec)\n",
      "INFO:tensorflow:global_step/sec: 21.4385\n",
      "INFO:tensorflow:loss = 0.3744754, step = 701 (4.664 sec)\n",
      "INFO:tensorflow:global_step/sec: 21.4136\n",
      "INFO:tensorflow:loss = 0.32597324, step = 801 (4.670 sec)\n",
      "INFO:tensorflow:global_step/sec: 22.1667\n",
      "INFO:tensorflow:loss = 0.27181357, step = 901 (4.511 sec)\n",
      "INFO:tensorflow:global_step/sec: 22.2173\n",
      "INFO:tensorflow:loss = 0.27270153, step = 1001 (4.501 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.5161\n",
      "INFO:tensorflow:loss = 0.26464644, step = 1101 (4.252 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.7369\n",
      "INFO:tensorflow:loss = 0.22492202, step = 1201 (4.213 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.562\n",
      "INFO:tensorflow:loss = 0.18850948, step = 1301 (4.244 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.7109\n",
      "INFO:tensorflow:loss = 0.1699378, step = 1401 (4.217 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.6578\n",
      "INFO:tensorflow:loss = 0.1903371, step = 1501 (4.227 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.4251\n",
      "INFO:tensorflow:loss = 0.12765962, step = 1601 (4.269 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.6546\n",
      "INFO:tensorflow:loss = 0.14119695, step = 1701 (4.227 sec)\n",
      "INFO:tensorflow:global_step/sec: 22.533\n",
      "INFO:tensorflow:loss = 0.121392526, step = 1801 (4.438 sec)\n",
      "INFO:tensorflow:global_step/sec: 21.126\n",
      "INFO:tensorflow:loss = 0.14298849, step = 1901 (4.734 sec)\n",
      "INFO:tensorflow:global_step/sec: 20.7395\n",
      "INFO:tensorflow:loss = 0.12086581, step = 2001 (4.821 sec)\n",
      "INFO:tensorflow:global_step/sec: 22.7284\n",
      "INFO:tensorflow:loss = 0.13730048, step = 2101 (4.400 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.1765\n",
      "INFO:tensorflow:loss = 0.12401772, step = 2201 (4.315 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.2306\n",
      "INFO:tensorflow:loss = 0.08946254, step = 2301 (4.305 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.2479\n",
      "INFO:tensorflow:loss = 0.09863055, step = 2401 (4.302 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.3789\n",
      "INFO:tensorflow:loss = 0.117545925, step = 2501 (4.277 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.8363\n",
      "INFO:tensorflow:loss = 0.08916695, step = 2601 (4.195 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.2566\n",
      "INFO:tensorflow:loss = 0.09194853, step = 2701 (4.300 sec)\n",
      "INFO:tensorflow:global_step/sec: 22.0659\n",
      "INFO:tensorflow:loss = 0.089303456, step = 2801 (4.532 sec)\n",
      "INFO:tensorflow:global_step/sec: 21.834\n",
      "INFO:tensorflow:loss = 0.0850828, step = 2901 (4.580 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.0115\n",
      "INFO:tensorflow:loss = 0.07665578, step = 3001 (4.346 sec)\n",
      "INFO:tensorflow:global_step/sec: 21.9299\n",
      "INFO:tensorflow:loss = 0.08190661, step = 3101 (4.560 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.4749\n",
      "INFO:tensorflow:loss = 0.05698895, step = 3201 (4.259 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.7282\n",
      "INFO:tensorflow:loss = 0.081066124, step = 3301 (4.214 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.5757\n",
      "INFO:tensorflow:loss = 0.05936368, step = 3401 (4.242 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.6944\n",
      "INFO:tensorflow:loss = 0.062697105, step = 3501 (4.220 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.0329\n",
      "INFO:tensorflow:loss = 0.055080947, step = 3601 (4.342 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.0493\n",
      "INFO:tensorflow:loss = 0.049559984, step = 3701 (4.339 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.0351\n",
      "INFO:tensorflow:loss = 0.055844955, step = 3801 (4.341 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.3941\n",
      "INFO:tensorflow:loss = 0.0505348, step = 3901 (4.274 sec)\n",
      "INFO:tensorflow:global_step/sec: 22.5813\n",
      "INFO:tensorflow:loss = 0.0520119, step = 4001 (4.429 sec)\n",
      "INFO:tensorflow:global_step/sec: 22.722\n",
      "INFO:tensorflow:loss = 0.049626835, step = 4101 (4.401 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.5911\n",
      "INFO:tensorflow:loss = 0.08396989, step = 4201 (4.239 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.2237\n",
      "INFO:tensorflow:loss = 0.04783809, step = 4301 (4.306 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.399\n",
      "INFO:tensorflow:loss = 0.050422754, step = 4401 (4.274 sec)\n",
      "INFO:tensorflow:global_step/sec: 22.669\n",
      "INFO:tensorflow:loss = 0.05665383, step = 4501 (4.411 sec)\n",
      "INFO:tensorflow:global_step/sec: 21.3909\n",
      "INFO:tensorflow:loss = 0.07028015, step = 4601 (4.675 sec)\n",
      "INFO:tensorflow:global_step/sec: 21.8558\n",
      "INFO:tensorflow:loss = 0.051450312, step = 4701 (4.575 sec)\n",
      "INFO:tensorflow:global_step/sec: 21.9794\n",
      "INFO:tensorflow:loss = 0.05008777, step = 4801 (4.550 sec)\n",
      "INFO:tensorflow:global_step/sec: 21.0022\n",
      "INFO:tensorflow:loss = 0.0463827, step = 4901 (4.761 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.7195\n",
      "INFO:tensorflow:loss = 0.027224278, step = 5001 (4.216 sec)\n",
      "INFO:tensorflow:global_step/sec: 24.1574\n",
      "INFO:tensorflow:loss = 0.045584396, step = 5101 (4.139 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.8324\n",
      "INFO:tensorflow:loss = 0.030244542, step = 5201 (4.196 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.9204\n",
      "INFO:tensorflow:loss = 0.0402273, step = 5301 (4.181 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.6923\n",
      "INFO:tensorflow:loss = 0.037832193, step = 5401 (4.221 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.8119\n",
      "INFO:tensorflow:loss = 0.03524958, step = 5501 (4.200 sec)\n",
      "INFO:tensorflow:global_step/sec: 24.0117\n",
      "INFO:tensorflow:loss = 0.0377734, step = 5601 (4.164 sec)\n",
      "INFO:tensorflow:global_step/sec: 24.3244\n",
      "INFO:tensorflow:loss = 0.05635477, step = 5701 (4.111 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.609\n",
      "INFO:tensorflow:loss = 0.02143224, step = 5801 (4.236 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.6933\n",
      "INFO:tensorflow:loss = 0.01922557, step = 5901 (4.220 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.6422\n",
      "INFO:tensorflow:loss = 0.039026648, step = 6001 (4.230 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.003\n",
      "INFO:tensorflow:loss = 0.031213952, step = 6101 (4.347 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.0101\n",
      "INFO:tensorflow:loss = 0.027385442, step = 6201 (4.346 sec)\n",
      "INFO:tensorflow:global_step/sec: 21.432\n",
      "INFO:tensorflow:loss = 0.030458279, step = 6301 (4.666 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.218\n",
      "INFO:tensorflow:loss = 0.03369922, step = 6401 (4.307 sec)\n",
      "INFO:tensorflow:global_step/sec: 21.7293\n",
      "INFO:tensorflow:loss = 0.03300318, step = 6501 (4.602 sec)\n",
      "INFO:tensorflow:global_step/sec: 22.5461\n",
      "INFO:tensorflow:loss = 0.027033214, step = 6601 (4.435 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.2169\n",
      "INFO:tensorflow:loss = 0.020006763, step = 6701 (4.308 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.1102\n",
      "INFO:tensorflow:loss = 0.023138283, step = 6801 (4.327 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.721\n",
      "INFO:tensorflow:loss = 0.021574501, step = 6901 (4.216 sec)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:global_step/sec: 22.8749\n",
      "INFO:tensorflow:loss = 0.040482827, step = 7001 (4.372 sec)\n",
      "INFO:tensorflow:global_step/sec: 22.3728\n",
      "INFO:tensorflow:loss = 0.031408735, step = 7101 (4.470 sec)\n",
      "INFO:tensorflow:global_step/sec: 22.4589\n",
      "INFO:tensorflow:loss = 0.019708859, step = 7201 (4.452 sec)\n",
      "INFO:tensorflow:global_step/sec: 21.6498\n",
      "INFO:tensorflow:loss = 0.027383469, step = 7301 (4.619 sec)\n",
      "INFO:tensorflow:global_step/sec: 22.3148\n",
      "INFO:tensorflow:loss = 0.023360977, step = 7401 (4.481 sec)\n",
      "INFO:tensorflow:global_step/sec: 22.1823\n",
      "INFO:tensorflow:loss = 0.026577486, step = 7501 (4.509 sec)\n",
      "INFO:tensorflow:global_step/sec: 22.018\n",
      "INFO:tensorflow:loss = 0.040591024, step = 7601 (4.541 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.7487\n",
      "INFO:tensorflow:loss = 0.020000473, step = 7701 (4.211 sec)\n",
      "INFO:tensorflow:global_step/sec: 23.8065\n",
      "INFO:tensorflow:loss = 0.03208575, step = 7801 (4.201 sec)\n",
      "INFO:tensorflow:Saving checkpoints for 7813 into /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpxxtascq7/model.ckpt.\n",
      "INFO:tensorflow:Loss for final step: 0.03523616.\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Restoring parameters from /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpxxtascq7/model.ckpt-7813\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "\n",
      "IN: apple\n",
      "OUT: a e l p p <EOS>\n",
      "\n",
      "IN: common\n",
      "OUT: c m m o o o <EOS>\n",
      "\n",
      "IN: zhedong\n",
      "OUT: d e g h n o z <EOS>\n"
     ]
    }
   ],
   "source": [
    "def infe_inps(str_li):\n",
    "    max_len = max([len(s) for s in str_li])\n",
    "    x_inps, x_seq_lens = [], []\n",
    "    for s in str_li:\n",
    "        x = [PARAMS['src_char2idx'].get(c, 3) for c in s] + [2]\n",
    "        x_inps.append(x)\n",
    "        x_seq_lens.append(len(x))\n",
    "    return {'src_idx': tf.keras.preprocessing.sequence.pad_sequences(x_inps, PARAMS['max_len'],\n",
    "                                                                     padding='post'),\n",
    "            'src_seq_lens': np.array(x_seq_lens)}\n",
    "\n",
    "\n",
    "def demo(xs, preds):\n",
    "    for x, pred in zip(xs, preds):\n",
    "        print('\\nIN: {}'.format(x))\n",
    "        x = np.array([PARAMS['src_char2idx'].get(c, 3) for c in x] + [2])\n",
    "        pred = x[pred[:len(x)]]\n",
    "        print('OUT: {}'.format(' '.join([PARAMS['src_idx2char'][i] for i in pred])))\n",
    "    \n",
    "\n",
    "def main():\n",
    "    src_idx, tgt_idx, src_seq_lens, tgt_seq_lens = preprocess_data()\n",
    "    \n",
    "    test_strs = ['apple', 'common', 'zhedong']\n",
    "    \n",
    "    estimator = tf.estimator.Estimator(model_fn)\n",
    "    \n",
    "    estimator.train(tf.estimator.inputs.numpy_input_fn(\n",
    "        x = {'src_idx': src_idx, 'src_seq_lens': src_seq_lens},\n",
    "        y = {'tgt_idx': tgt_idx, 'tgt_seq_lens': tgt_seq_lens},\n",
    "        batch_size = PARAMS['batch_size'],\n",
    "        num_epochs = PARAMS['n_epochs'],\n",
    "        shuffle = True))\n",
    "    \n",
    "    preds = list(estimator.predict(tf.estimator.inputs.numpy_input_fn(\n",
    "        x = infe_inps(test_strs),\n",
    "        shuffle = False)))\n",
    "    \n",
    "    demo(test_strs, preds)\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
